
import os
import cv2
from pycocotools.coco import COCO


DISPLAY_TIME = 0
# DISPLAY_TIME = 100

DISPLAY_WIDTH = 1920
DISPLAY_HEIGHT = 1080

display_class = True
skip_crowd = True


CLASSES = ['person', 'cat', 'dog', 'baby_carriage', 'face', 'cell phone', 'bicycle', 'motorcycle']
# CLASSES = ['person']
# CLASSES = ['cat']
# CLASSES = ['dog']
# CLASSES = ['cell phone']
# CLASSES = ['bicycle']
# CLASSES = ['motorcycle']


ROOT_PATH = r'/rootfs/media/yery/Kaso/data'

# USE_DATASET = 'BDD100K/train'
USE_DATASET = 'BDD100K/val'


COCO_DATASET_INFO = {
    'BDD100K': {
        'train': {
            'IMAGE_FOLDER': r'BDD100K/out/Images/100k/train',
            'ANNO_FOLDER': r'BDD100K/out',
            'ANNO_FILENAME': r'bdd100k_coco_labels_images_train.json',
        },
        'val': {
            'IMAGE_FOLDER': r'BDD100K/out/Images/100k/val',
            'ANNO_FOLDER': r'BDD100K/out',
            'ANNO_FILENAME': r'bdd100k_coco_labels_images_val.json',
        },
    },
}


CATEGORIES = [
    {
        'id': 1,
        'name': 'person',
        'supercategory': 'person',  # person
    },
    {
        'id': 17,
        'name': 'cat',
        'supercategory': 'cat',  # animal
    },
    {
        'id': 18,
        'name': 'dog',
        'supercategory': 'dog',  # animal
    },
    {
        'id': 102,
        'name': 'baby_carriage',
        'supercategory': 'baby_carriage',
    },
    {
        'id': 103,
        'name': 'face',
        'supercategory': 'face',  # person
    },
    {
        'id': 77,
        'name': 'cell phone',
        'supercategory': 'cell phone',  # electronic
    },
    {
        'id': 2,
        'name': 'bicycle',
        'supercategory': 'bicycle',  # vehicle
    },
    {
        'id': 4,
        'name': 'motorcycle',
        'supercategory': 'motorcycle',  # vehicle
    },
]

coco_id_lables = {category['id']: category['name'] for category in CATEGORIES}

# CATEGORIES = [
#     {"supercategory": "person","id": 1,"name": "person"},{"supercategory": "vehicle","id": 2,"name": "bicycle"},{"supercategory": "vehicle","id": 3,"name": "car"},{"supercategory": "vehicle","id": 4,"name": "motorcycle"},{"supercategory": "vehicle","id": 5,"name": "airplane"},{"supercategory": "vehicle","id": 6,"name": "bus"},{"supercategory": "vehicle","id": 7,"name": "train"},{"supercategory": "vehicle","id": 8,"name": "truck"},{"supercategory": "vehicle","id": 9,"name": "boat"},{"supercategory": "outdoor","id": 10,"name": "traffic light"},{"supercategory": "outdoor","id": 11,"name": "fire hydrant"},{"supercategory": "outdoor","id": 13,"name": "stop sign"},{"supercategory": "outdoor","id": 14,"name": "parking meter"},{"supercategory": "outdoor","id": 15,"name": "bench"},{"supercategory": "animal","id": 16,"name": "bird"},{"supercategory": "animal","id": 17,"name": "cat"},{"supercategory": "animal","id": 18,"name": "dog"},{"supercategory": "animal","id": 19,"name": "horse"},{"supercategory": "animal","id": 20,"name": "sheep"},{"supercategory": "animal","id": 21,"name": "cow"},{"supercategory": "animal","id": 22,"name": "elephant"},{"supercategory": "animal","id": 23,"name": "bear"},{"supercategory": "animal","id": 24,"name": "zebra"},{"supercategory": "animal","id": 25,"name": "giraffe"},{"supercategory": "accessory","id": 27,"name": "backpack"},{"supercategory": "accessory","id": 28,"name": "umbrella"},{"supercategory": "accessory","id": 31,"name": "handbag"},{"supercategory": "accessory","id": 32,"name": "tie"},{"supercategory": "accessory","id": 33,"name": "suitcase"},{"supercategory": "sports","id": 34,"name": "frisbee"},{"supercategory": "sports","id": 35,"name": "skis"},{"supercategory": "sports","id": 36,"name": "snowboard"},{"supercategory": "sports","id": 37,"name": "sports ball"},{"supercategory": "sports","id": 38,"name": "kite"},{"supercategory": "sports","id": 39,"name": "baseball bat"},{"supercategory": "sports","id": 40,"name": "baseball glove"},{"supercategory": "sports","id": 41,"name": "skateboard"},{"supercategory": "sports","id": 42,"name": "surfboard"},{"supercategory": "sports","id": 43,"name": "tennis racket"},{"supercategory": "kitchen","id": 44,"name": "bottle"},{"supercategory": "kitchen","id": 46,"name": "wine glass"},{"supercategory": "kitchen","id": 47,"name": "cup"},{"supercategory": "kitchen","id": 48,"name": "fork"},{"supercategory": "kitchen","id": 49,"name": "knife"},{"supercategory": "kitchen","id": 50,"name": "spoon"},{"supercategory": "kitchen","id": 51,"name": "bowl"},{"supercategory": "food","id": 52,"name": "banana"},{"supercategory": "food","id": 53,"name": "apple"},{"supercategory": "food","id": 54,"name": "sandwich"},{"supercategory": "food","id": 55,"name": "orange"},{"supercategory": "food","id": 56,"name": "broccoli"},{"supercategory": "food","id": 57,"name": "carrot"},{"supercategory": "food","id": 58,"name": "hot dog"},{"supercategory": "food","id": 59,"name": "pizza"},{"supercategory": "food","id": 60,"name": "donut"},{"supercategory": "food","id": 61,"name": "cake"},{"supercategory": "furniture","id": 62,"name": "chair"},{"supercategory": "furniture","id": 63,"name": "couch"},{"supercategory": "furniture","id": 64,"name": "potted plant"},{"supercategory": "furniture","id": 65,"name": "bed"},{"supercategory": "furniture","id": 67,"name": "dining table"},{"supercategory": "furniture","id": 70,"name": "toilet"},{"supercategory": "electronic","id": 72,"name": "tv"},{"supercategory": "electronic","id": 73,"name": "laptop"},{"supercategory": "electronic","id": 74,"name": "mouse"},{"supercategory": "electronic","id": 75,"name": "remote"},{"supercategory": "electronic","id": 76,"name": "keyboard"},{"supercategory": "electronic","id": 77,"name": "cell phone"},{"supercategory": "appliance","id": 78,"name": "microwave"},{"supercategory": "appliance","id": 79,"name": "oven"},{"supercategory": "appliance","id": 80,"name": "toaster"},{"supercategory": "appliance","id": 81,"name": "sink"},{"supercategory": "appliance","id": 82,"name": "refrigerator"},{"supercategory": "indoor","id": 84,"name": "book"},{"supercategory": "indoor","id": 85,"name": "clock"},{"supercategory": "indoor","id": 86,"name": "vase"},{"supercategory": "indoor","id": 87,"name": "scissors"},{"supercategory": "indoor","id": 88,"name": "teddy bear"},{"supercategory": "indoor","id": 89,"name": "hair drier"},{"supercategory": "indoor","id": 90,"name": "toothbrush"}
# ]


def coco_select_classes_display(image_folder, anno_folder, anno_filename, classes_list):
    anno_file = os.path.join(anno_folder, anno_filename)
    coco = COCO(anno_file)

    # get all images containing given categories, select one at random
    catIds = coco.getCatIds(catNms=classes_list)
    _img_ids = set()
    for cat_id in catIds:
        _img_ids |= set(coco.getImgIds(catIds=cat_id))
    img_ids = list(_img_ids)

    objCntPreCls = {}
    for id in catIds:
        objCntPreCls[id] = 0

    file_count = 0
    for i in img_ids:
        file_count = file_count + 1
        if (file_count % 100) == 0:
            print('Proccess', file_count, 'Files.')

        image_info = coco.loadImgs([i])[0]

        img_id = image_info['id']
        ann_ids = coco.getAnnIds(imgIds=[img_id])
        ann_info = coco.loadAnns(ann_ids)
        has_select_classes = False
        label_box_list = []
        for j, ann in enumerate(ann_info):
            if ann['category_id'] not in catIds:
                continue
            if skip_crowd:
                if ann["iscrowd"]:
                    continue
            label = coco_id_lables[ann['category_id']]
            label_box_list.append([label, ann['bbox']])
            # "area": area
            has_select_classes = True
            objCntPreCls[ann['category_id']] += 1

        if has_select_classes:
            image_path = os.path.join(image_folder, image_info['file_name'])
            image = cv2.imread(image_path)
            for lable, box in label_box_list:
                cv2.rectangle(image, (int(box[0]), int(box[1])), (int(box[0]+box[2]), int(box[1]+box[3])), (0, 255, 0), thickness=1)
                if display_class:
                    cv2.putText(image, lable, (int(box[0])+2, int(box[1]) + 16), cv2.FONT_HERSHEY_COMPLEX, 0.5, (0, 0, 255))
            # width = image.shape[1]
            # height = image.shape[0]
            # sw = DISPLAY_WIDTH / width
            # sh = DISPLAY_HEIGHT / height
            # scale = min(sh, sw)
            # width = int(round(width*scale))
            # height = int(round(height * scale))
            # image = cv2.resize(image, (width, height))
            cv2.imshow('image', image)
            if DISPLAY_TIME < 1:
                key = cv2.waitKey()
            else:
                key = cv2.waitKey(DISPLAY_TIME)
            if key == 27 or key == ord('q') or key == ord('Q'):
                break
    cv2.destroyAllWindows()
    print('Proccess', file_count, 'Files.')

    for id, objCnt in objCntPreCls.items():
        for c in CATEGORIES:
            if c['id'] == id:
                print(c['name'], objCnt)
                break


if __name__ == "__main__":
    DATASET, STAGE = USE_DATASET.split('/')
    IMAGE_FOLDER = os.path.join(ROOT_PATH, COCO_DATASET_INFO[DATASET][STAGE]['IMAGE_FOLDER'])
    ANNO_FOLDER = os.path.join(ROOT_PATH, COCO_DATASET_INFO[DATASET][STAGE]['ANNO_FOLDER'])
    ANNO_FILENAME = COCO_DATASET_INFO[DATASET][STAGE]['ANNO_FILENAME']
    coco_select_classes_display(IMAGE_FOLDER, ANNO_FOLDER, ANNO_FILENAME, CLASSES)
